import pysam
import numpy as np
import pandas as pd
import argparse
from collections import defaultdict

from deepsignal3.utils import bam_reader


def get_q2tloc_from_cigar(r_cigar_tuple, strand, seq_len):
    """
    insertion: -1, deletion: -2, mismatch: -3
    :param r_cigar_tuple: pysam.alignmentSegment.cigartuples
    :param strand: 1/-1 for fwd/rev
    :param seq_len: read alignment length
    :return: query pos to ref pos
    """
    fill_invalid = -2
    # get each base calls genomic position
    q_to_r_poss = np.full(seq_len + 1, fill_invalid, dtype=np.int32)
    # process cigar ops in read direction
    curr_r_pos, curr_q_pos = 0, 0
    cigar_ops = r_cigar_tuple if strand == 1 else r_cigar_tuple[::-1]
    for op, op_len in cigar_ops:
        if op == 1:
            # inserted bases into ref
            for q_pos in range(curr_q_pos, curr_q_pos + op_len):
                q_to_r_poss[q_pos] = -1
            curr_q_pos += op_len
        elif op in (2, 3):
            # deleted ref bases
            curr_r_pos += op_len
        elif op in (0, 7, 8):
            # aligned bases
            for op_offset in range(op_len):
                q_to_r_poss[curr_q_pos + op_offset] = curr_r_pos + op_offset
            curr_q_pos += op_len
            curr_r_pos += op_len
        elif op == 6:
            # padding (shouldn't happen in mappy)
            pass
    q_to_r_poss[curr_q_pos] = curr_r_pos
    if q_to_r_poss[-1] == fill_invalid:
        raise ValueError(
            (
                "Invalid cigar string encountered. Reference length: {}  Cigar "
                + "implied reference length: {}"
            ).format(seq_len, curr_r_pos)
        )
    return q_to_r_poss

def process_data(seq_read):
    """
    处理每一个读取，提取CG位点的染色体和位置
    """
    features_list = []
    seq = seq_read.get_forward_sequence()
    
    # 只处理未映射的reads
    if seq_read.is_unmapped or seq is None:
        return features_list

    # 获取染色体和方向信息
    ref_name = seq_read.reference_name
    strand = "-" if seq_read.is_reverse else "+"
    strand_code = -1 if seq_read.is_reverse else 1
    ref_start = seq_read.reference_start
    ref_end = seq_read.reference_end
    cigar_tuples = seq_read.cigartuples
    qalign_start = seq_read.query_alignment_start
    qalign_end = seq_read.query_alignment_end
    if seq_read.is_reverse:
            seq_start = len(seq) - qalign_end
            seq_end = len(seq) - qalign_start
    else:
        seq_start = qalign_start
        seq_end = qalign_end
    q_to_r_poss = get_q2tloc_from_cigar(
        cigar_tuples, strand_code, (seq_end - seq_start)
    )
    # 获取CG位点位置
    cg_locs = [i for i in range(len(seq) - 1) if seq[i:i+2] == "CG"]
    kmer_len=5
    num_bases = (kmer_len - 1) // 2

    for loc_in_read in cg_locs:
        k_mer = seq[(loc_in_read - num_bases):(loc_in_read + num_bases + 1)]
        if loc_in_read + num_bases + 1 >= len(seq) or loc_in_read - num_bases < 0:
            continue
        if seq_start <= loc_in_read < seq_end:
            offset_idx = loc_in_read - seq_start
            if q_to_r_poss[offset_idx] != -1:
                if strand == "-":
                    # pos = '.'#loc_in_read
                    ref_pos = ref_end - 1 - q_to_r_poss[offset_idx]
                else:
                    # pos = loc_in_read
                    ref_pos = ref_start + q_to_r_poss[offset_idx]
            else:
                continue
        else:
            continue

        features_list.append((ref_name, ref_pos, strand, k_mer))
    
    return features_list
def assign_keys_to_bins(bed_data, bins, labels):
    """将甲基化水平分配到相应的区间"""
    binned_positions = {label: set() for label in labels}
    for chrom, pos, methylation_level in bed_data:
        if methylation_level == 0:
            binned_positions['[0,0]'].add((chrom, pos))
        elif methylation_level == 100:
            binned_positions['[100,100]'].add((chrom, pos))
        else:
            for i in range(len(bins) - 1):
                if bins[i] < methylation_level <= bins[i + 1]:
                    binned_positions[labels[i + 1]].add((chrom, pos))
                    break
    return binned_positions
def main(args):
    # 读取BED文件
    bed_file = args.bed_file
    #bed_df = pd.read_csv(bed_file, sep='\t', header=None, usecols=[0, 1, 10], names=['chrom', 'pos', 'methylation_level'])
    bed_data = []
    with open(bed_file, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            chrom = parts[0]
            pos = int(parts[1])
            methylation_level = float(parts[10])
            bed_data.append((chrom, pos, methylation_level))
    # 将甲基化水平划分到不同区间
    bins = [0, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
    labels = ['[0,0]', '(0,10]', '(10,20]', '(20,30]', '(30,40]', '(40,50]', '(50,60]', 
              '(60,70]', '(70,80]', '(80,90]', '(90,100)', '[100,100]']
    binned_positions = assign_keys_to_bins(bed_data, bins, labels)
    # intervals = {
    #     '[0,0]': (0, 0),
    #     '(0,10]': (0, 10),
    #     '(10,20]': (10, 20),
    #     '(20,30]': (20, 30),
    #     '(30,40]': (30, 40),
    #     '(40,50]': (40, 50),
    #     '(50,60]': (50, 60),
    #     '(60,70]': (60, 70),
    #     '(70,80]': (70, 80),
    #     '(80,90]': (80, 90),
    #     '(90,100)': (90, 100),
    #     '[100,100]': (100, 100)
    # }

    # # 将位点分配到不同的区间
    # bed_df['interval'] = bed_df['methylation_level'].apply(
    #     lambda x: next((k for k, v in intervals.items() if v[0] <= x <= v[1]), None)
    # )

    # 读取BAM文件
    bam_file = args.bam_file
    bam = bam_reader.ReadIndexedBam(bam_file)
    bam.open()

    # 统计每个区间内的CG位点频率
    frequency = defaultdict(lambda: defaultdict(int))

    # 遍历BAM文件
    for read in bam:
        features = process_data(read)
        if len(features) == 0:
            continue
        for ref_name, ref_pos, strand, k_mer in features:
            for interval, positions in binned_positions.items():
                if (ref_name, ref_pos) in positions:
                    frequency[interval][k_mer] += 1
                    break
        #     # 根据BED文件中的甲基化水平查找对应的区间
        #     methylation_level = bed_df.loc[
        #         (bed_df['chrom'] == ref_name) & (bed_df['pos'] == ref_pos), 'methylation_level'
        #     ]
        #     if not methylation_level.empty:
        #         interval = next(
        #             (k for k, v in intervals.items() if v[0] <= methylation_level.values[0] <= v[1]),
        #             None
        #         )
        #         if interval:
        #             # 统计5mer在相应区间的频率
        #             frequency[interval][k_mer] += 1

    # 关闭BAM文件
    bam.close()

    # 将统计结果写入输出文件
    with open(args.output_file, 'w') as output:
        for interval, kmer_dict in frequency.items():
            output.write(f"Interval: {interval}\n")
            for kmer, count in kmer_dict.items():
                output.write(f"{kmer}: {count}\n")
            output.write("\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="统计BAM文件中各甲基化区间的CG位点频率")
    parser.add_argument("bam_file", help="输入BAM文件路径")
    parser.add_argument("bed_file", help="输入BED文件路径")
    parser.add_argument("output_file", help="输出结果文件路径")

    args = parser.parse_args()
    main(args)
